//==============================================
//file info//===================================
//file   :test_paillier_gmp.c
//author :王晓晨 HUST beamworld.gitee.io
//date   :2023.05.28
//change :
//   1)完成paillier密钥生成和加解密 2023.05.28
//=============================================

//==============================================
//library//=====================================
#include <gmp.h>
#include <test_paillier_gmp.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
//==============================================

//==============================================
//define//======================================
//调试信息部分-----------------------------------
//#define DEBUG_LOG			//调试信息开关
//#define DEBUG_RANDOM		//随机数生成函数调试信息开关
//#define DEBUG_GENERATE	//密钥生成函数调试信息开关
#define DEBUG_ENCRYPT		//加密函数调试信息开关
//#define DEBUG_DECRYPT		//解密函数调试信息开关
//----------------------------------------------

//参数定义部分-----------------------------------
#define KEYPAIR_BIT_NUM		1024	//密钥位数
#define PLAINTEXT_BIT_NUM 	64		//明文位数
#define INPUT_MESSAGE 		"18446744073709551615"	//明文内容,位数不大于 PLAINTEXT_BIT_NUM, 无符号下0 ~ 18446744073709551615(20位)
//----------------------------------------------
//==============================================

//==============================================
//main//========================================
void main(void)
{
	//输出启动信息
    printf("\n@>>>program starting========================================================\n");

    paillier_keypair* keypair = paillier_keypair_new();	//初始化密钥对变量
    generate_keypair(keypair, KEYPAIR_BIT_NUM);	//获取密钥对,长度为 KEYPAIR_BIT_NUM

	mpz_t input, cipher_text, plain_text;	//明文，密文，解密文
	mpz_inits(input, cipher_text, plain_text, NULL);	//初始化为NULL

	mpz_set_str(input, INPUT_MESSAGE, 10);	//设置明文信息
	gmp_printf("明文\t:%Zd\n", input);	//输出明文信息

	encrypt(cipher_text, keypair->pubkey, input);	//加密信息
	gmp_printf("密文\t:%Zd\n", cipher_text);	//输出密文信息

	decrypt(plain_text, keypair, cipher_text);	//解密信息
	gmp_printf("解密文\t:%Zd\n", plain_text);	//输出解密文信息
	
	//清理
	mpz_clears(plain_text, cipher_text, input, NULL);
	paillier_keypair_free(keypair);
	printf("@>>>program terminated======================================================\n");
}
//==============================================

//==============================================
//function//====================================
//随机数生产函数部分-----------------------------
/*
 * [函数名称] generate_urandomm
 * [功能描述] 生成0 ~ n - 1的随机无符整数
 * [参数描述] rop 随机数结果
 *			 n 随机数上限
 * [ 返回值 ] void	
 */
void generate_urandomm(mpz_t rop, mpz_t n) 
{	
	struct timeval tv;  //时间变量，用于随机数种子
	gmp_randstate_t random_state;	//随机数状态类型，没有深究

	gettimeofday(&tv, NULL);    //获得精确时间
    #ifdef DEBUG_RANDOM
        printf("@>>>get system time\n");
        printf("time = %ld\n", tv.tv_sec);    //调试输出精确时间
    #endif

    gmp_randinit_default(random_state);    //将随机数状态设为默认
    gmp_randseed_ui(random_state, tv.tv_sec);  //精确时间作为随机数种子,这里应是随机数不随机的关键

	mpz_init(rop);	//初始化为0
	mpz_urandomm(rop, random_state, n);	//生成介于 0 到 n-1 的随机整数
	gmp_randclear(random_state);	//清理随机数
}

/*
 * [函数名称] generate_urandomb
 * [功能描述] 生成介于 0 到 2^k（k 为正整数）之间的随机无符整数
 * [参数描述] rop 随机数结果
 *			 k 随机数上限
 * [ 返回值 ] void	
 */
void generate_urandomb(mpz_t rop, mp_bitcnt_t k) 
{
	struct timeval tv;  //时间变量，用于随机数种子
	gmp_randstate_t random_state;	//随机数状态类型，没有深究

	gettimeofday(&tv, NULL);    //获得精确时间
    #ifdef DEBUG_RANDOM
        printf("@>>>get system time\n");
        printf("time = %ld\n", tv.tv_sec);    //调试输出精确时间
    #endif

    gmp_randinit_default(random_state);    //将随机数状态设为默认
    gmp_randseed_ui(random_state, tv.tv_sec);  //精确时间作为随机数种子,这里应是随机数不随机的关键

	mpz_init(rop);	//初始化为0
	mpz_urandomb(rop, random_state, k);	//生成介于 0 到 2^k（k 为正整数）之间的随机整数
	gmp_randclear(random_state);	//清理随机数
}
//----------------------------------------------

//paillier密钥生成，加解密部分--------------------
/*
 * [函数名称] generate_keypair
 * [功能描述] 生成密钥
 * [参数描述] p_kp 密钥结构体
 *			 bits 密钥位数
 * [ 返回值 ] void	
 */
void generate_keypair(paillier_keypair* p_kp, int bits)
{
	#ifdef DEBUG_ENCRYPT
        printf("@>>>generate_keypair function\n");
    #endif

    mpz_t p, q;    		//随机大质数p, q
	mpz_t p_subtract_1;	//p-1
	mpz_t q_subtract_1;	//q-1
	mpz_t x;			//随机数x，x属于Zn*
	mpz_t h;			//过程量h
	mpz_t hs;			//公钥hs
	mpz_t n;			//公钥n
	mpz_t n_square;		//n^2
	mpz_t lambda;		//私钥lambda
	mpz_t mu;			//私钥mu
	mpz_t temp;    		//临时变量

    mpz_init(p);    	//初始化p为0
	mpz_init(q);    	//初始化q为0
	mpz_init(p_subtract_1);	//初始化p_subtract_1为0
	mpz_init(q_subtract_1);	//初始化q_subtract_1为0
	mpz_init(x);    	//初始化z为0
	mpz_init(h);    	//初始化h为0
	mpz_init(hs);   	//初始化hs为0
	mpz_init(n);    	//初始化n为0
	mpz_init(n_square);	//n_square为0
	mpz_init(lambda);   //初始化lambda为0
	mpz_init(mu);    	//初始化mu为0
	mpz_init(temp);		//初始化temp为0

	p_kp->bits = bits;			//密钥对位数
	p_kp->pubkey->bits = bits;	//公钥位数
	p_kp->prikey->bits = bits;	//私钥位数

	//令p q 中的p大于要加密文的最大位数所能代表的最大数，这样可以保证明文m一定小于n
	mpz_init_set_ui(temp, 2);	//temp取2
	mpz_pow_ui(temp, temp, PLAINTEXT_BIT_NUM / 2);	//计算q最小位数为2^(明文最大位数 / 2)所能代表的最大数 + 1

    for(generate_urandomb(p, bits / 2); mpz_cmp(p, temp) < 0; mpz_add_ui(p, p, 1))	//q不满足位数则加一
	//generate_urandomb(p, bits / 2);  //为p设置随机数,范围 0 ~ 2^bits
    mpz_nextprime(p, p);    //p设置为p临近大质数
	//获取模4为3的质数p
	//mpz_mod_ui应该可以返回余数，可以进一步修改
	for(mpz_mod_ui(temp, p, 4); mpz_cmp_ui(temp, 3) != 0; mpz_mod_ui(temp, p, 4)) 
	{
		mpz_nextprime(p, p);	//p设置为p临近大质数
	}
    #ifdef DEBUG_GENERATE
        printf("@>>>set p with random (p为质数,模4为3)\n");
        gmp_printf("p = %Zd\n", p);   //调试输出p
    #endif

	for(mpz_init_set(q, p); mpz_cmp(q, p) == 0;)
	{
		generate_urandomb(q, bits / 2);  //为q设置随机数,范围 0 ~ RANDOM_MAX
		mpz_nextprime(q, q);	//q设置为q临近大质数
		for (mpz_mod_ui(temp, q, 4); mpz_cmp_ui(temp, 3) != 0; mpz_mod_ui(temp, q, 4)) 
		{
			mpz_nextprime(q, q);	//p设置为p临近大质数
		}
	}
    #ifdef DEBUG_GENERATE
        printf("@>>>set q with random (q为质数,模4为3)\n");
        gmp_printf("q = %Zd\n", q);   //调试输出q
    #endif

	mpz_mul(n, p, q);	//n=p*q，是公钥一部分,写入公钥
	mpz_pow_ui(n_square, n, 2);
	#ifdef DEBUG_GENERATE
        printf("@>>>set n = q * p, n_square = n^2\n");
        gmp_printf("n = %Zd , n_square = %Zd\n", n, n_square);   //调试输出n
    #endif

	mpz_sub_ui(p_subtract_1, p, 1);	//p_subtract_1 = (p - 1)
	mpz_sub_ui(q_subtract_1, q, 1);	//q_subtract_1 = (q - 1)
	mpz_mul(temp, p_subtract_1, q_subtract_1); // temp = (p - 1)(q - 1)
	mpz_tdiv_q_ui(lambda, temp, 2);	// lambda = (p - 1)(q - 1) / 2
	#ifdef DEBUG_GENERATE
        printf("@>>>set lambda = (q-1) * (p-1) / 2\n");
        gmp_printf("lambda = %Zd\n", lambda);   //调试输出lambda
    #endif

	//当g = n + 1时简化mu的计算公式为mu = lambda^-1 mod n
	mpz_invert(mu, lambda, n);
	#ifdef DEBUG_GENERATE
        printf("@>>>set mu = lambda^-1 mod n\n");
        gmp_printf("mu = %Zd\n", mu);   //调试输出x
    #endif

	generate_urandomm(x, n);	//生成随机数x, z属于0 ~ n-1
	if(mpz_cmp_ui(x, 0) == 0)	//x为0则加1
	{
		mpz_add_ui(x, x, 1);
	}
	#ifdef DEBUG_GENERATE
        printf("@>>>set x = 随机数, 范围 1 ~ n-1\n");
        gmp_printf("x = %Zd\n", x);   //调试输出x
    #endif

	mpz_pow_ui(h, x, 2);	//h = x^2
	mpz_neg(h, h);	//h = -h
	#ifdef DEBUG_GENERATE
        printf("@>>>set h = -x^2\n");
        gmp_printf("h = %Zd\n", h);   //调试输出h
    #endif

	mpz_powm(hs, h, n, n_square);	//hs = h^n mod n^2
	#ifdef DEBUG_GENERATE
        printf("@>>>set hs = h^n mod n^2\n");
        gmp_printf("hs = %Zd\n", hs);   //调试输出hs
    #endif

	//写入密钥
	mpz_set(p_kp->prikey->lambda, lambda);
	mpz_set(p_kp->prikey->mu, mu);
	mpz_set(p_kp->pubkey->hs, hs);
	mpz_set(p_kp->pubkey->n, n);
	mpz_set(p_kp->pubkey->n_square, n_square);

	#ifdef DEBUG_LOG
        printf("@>>>密钥写入完成\n");
		printf("@私钥========================\n");
        gmp_printf("&prikey.lambda = %Zd\n&prikey.mu = %Zd\n", p_kp->prikey->lambda, p_kp->prikey->mu);   //调试输出hs
		printf("@公钥========================\n");
		gmp_printf("&pubkey.n = %Zd\n&pubkey.hs = %Zd\n", p_kp->pubkey->n, p_kp->pubkey->hs);   //调试输出hs
		printf("@============================\n");
    #endif

	//清理
    mpz_clear(p);
	mpz_clear(q);
	mpz_clear(p_subtract_1);
	mpz_clear(q_subtract_1);
	mpz_clear(x);
	mpz_clear(h);
	mpz_clear(n);
	mpz_clear(n_square);
	mpz_clear(lambda);
	mpz_clear(mu);
	mpz_clear(temp);
}
//----------------------------------------------
/*
 * [函数名称] encrypt
 * [功能描述] 加密明文
 * [参数描述] pubkey 公钥结构体
 *			 out 密文
 *			 m 明文
 * [ 返回值 ] 1 密钥位数不正确 
 * 			 0 正常
 */
void encrypt(mpz_t out, paillier_pubkey* pubkey, mpz_t m) 
{
	#ifdef DEBUG_ENCRYPT
        	printf("@>>>encrypt function\n");
    #endif

	mpz_t temp_1, temp_2, nm, alpha;	//两个临时变量，n*m 和 随机数alpha
	mpz_inits(temp_1, temp_2, nm, alpha, NULL);	//初始化变量
	
	generate_urandomb(alpha, pubkey->bits / 2);	//生成随机数alpha 属于0 ~ 2^(bits/2)
	#ifdef DEBUG_ENCRYPT
        printf("@>>>set alpha = 0 ~ 2^(bits/2)\n");
        gmp_printf("alpha = %Zd\n", alpha);   //调试输出
    #endif

	mpz_mul(nm, pubkey->n, m);	// nm = n * m
	#ifdef DEBUG_ENCRYPT
        printf("@>>>set nm = n * m\n");
        gmp_printf("nm = %Zd\n", nm);   //调试输出
    #endif

	mpz_add_ui(temp_1, nm, 1);	// temp_1 = n * m + 1
	mpz_powm(temp_2, pubkey->hs, alpha, pubkey->n_square);	// temp_2 = hs^alpha mod n^2
	mpz_mul(out, temp_1, temp_2);	// out = (n * m + 1) * (hs^alpha)


	#ifdef DEBUG_LOG
        printf("@>>>明文加密完成,");
        gmp_printf("密文 = %Zd\n", out);   //调试输出
    #endif

	//清理
	mpz_clears(alpha, nm, temp_1, temp_2, NULL);
}

/*
 * [函数名称] decrypt
 * [功能描述] 解密密文
 * [参数描述] p_kp 密钥结构体
 *			 out 明文
 *			 c 密文
 * [ 返回值 ] void
 */
void decrypt(mpz_t out, paillier_keypair* p_kp, mpz_t c) 
{
	#ifdef DEBUG_DECRYPT
        	printf("@>>>decrypt function\n");
    #endif

	mpz_t l_x;	//l_x为函数L(x) = (x - 1) / n 的结果
	mpz_inits(l_x, NULL);

	mpz_powm(l_x, c, p_kp->prikey->lambda, p_kp->pubkey->n_square);	// l_x = c^lambda mod n^2 , 这里计算的是l_x的输入x,还没有计算L(x)函数
	mpz_sub_ui(l_x, l_x, 1);	// l_x = l_x - 1 ,计算L(x)函数分子
	mpz_tdiv_q(out, l_x, p_kp->pubkey->n); // out = c^lambda mod n^2 - 1 / n  , div函数q和r等不能使用同一个变量否则行为不定,故 out 代替l_x
	#ifdef DEBUG_DECRYPT
        printf("@>>>set L(x) = c^lambda mod n^2 - 1\n");
        gmp_printf("L(x) = %Zd\n", out);   //调试输出
    #endif

	mpz_mul(out, out, p_kp->prikey->mu);	//out = l_x * mu
	mpz_mod(out, out, p_kp->pubkey->n);	//out = ((c^lambda mod n^2 - 1) / n ) * mu mod n
	#ifdef DEBUG_LOG
        printf("@>>>解密完成,");
        gmp_printf("解密文 = %Zd\n", out);   //调试输出
    #endif

	//清理
	mpz_clears(l_x, NULL);
}
//----------------------------------------------

//paillier密钥结构体初始化和释放部分--------------
/*
 * [函数名称] paillier_keypair_new
 * [功能描述] 生成新的密钥结构体
 * [参数描述] NULL
 * [ 返回值 ] paillier_keypair* paillier密钥结构体指针
 */
paillier_keypair* paillier_keypair_new() 
{
	paillier_keypair* keypair = (paillier_keypair*)malloc(sizeof(paillier_keypair));
	if (keypair) {
		keypair->bits = 0;
		keypair->prikey = (paillier_prikey*)malloc(sizeof(paillier_prikey));
		if (keypair->prikey) {
			keypair->prikey->bits = 0;
			mpz_init(keypair->prikey->lambda);
			mpz_init(keypair->prikey->mu);
		}
		keypair->pubkey = (paillier_pubkey*)malloc(sizeof(paillier_pubkey));
		if (keypair->pubkey) {
			keypair->pubkey->bits = 0;
			mpz_init(keypair->pubkey->hs);
			mpz_init(keypair->pubkey->n);
			mpz_init(keypair->pubkey->n_square);
		}
	}
	return keypair;
}

/*
 * [函数名称] paillier_keypair_free
 * [功能描述] 释放密钥结构体
 * [参数描述] keypair 要释放的密钥结构体指针
 * [ 返回值 ] void
 */
void paillier_keypair_free(paillier_keypair* keypair) 
{
	free(keypair->prikey);
	free(keypair->pubkey);
	free(keypair);
	return;
}
//----------------------------------------------
//==============================================